import collections


# prim算法
class Solution(object):
    def minimumCost(self, n, connections):
        """
        :type n: int
        :type connections: List[List[int]]
        :rtype: int
        """
        count = 1
        edges = collections.defaultdict(dict)
        for item in connections:
            edges[item[0]][item[1]] = item[2]
            edges[item[1]][item[0]] = item[2]
        already_set = {1}
        value = 0
        while count < n:
            count += 1
            min_value = 10 ** 5 + 1
            min_id = - 1
            for item in already_set:
                for key in edges[item]:
                    if key not in already_set and edges[item][key] < min_value:
                        min_value = edges[item][key]
                        min_id = key
            if min_id == -1:
                return -1
            value += min_value
            already_set.add(min_id)
        return value


# cruscal算法
class Solution(object):
    def minimumCost(self, n, connections):
        """
        :type n: int
        :type connections: List[List[int]]
        :rtype: int
        """
        if n == 1:
            return 0
        p = [i for i in range(n + 1)]
        connections.sort(key=lambda x: x[2])
        self.part = n
        self.value = 0

        def find(x):
            if p[x] == x:
                return x
            else:
                p[x] = find(p[x])
                return p[x]

        def union(x, y, z):
            rootx = find(x)
            rooty = find(y)
            if rootx != rooty:
                self.part -= 1
                self.value += z
                if rootx < rooty:
                    p[rooty] = rootx
                else:
                    p[rootx] = rooty

        for item in connections:
            union(item[0], item[1], item[2])
            if self.part == 1:
                return self.value

        return -1


data = Solution()
n = 3
connections = [[1, 2, 5], [1, 3, 6], [2, 3, 1]]
print(data.minimumCost(n, connections))
n = 4
connections = [[1, 2, 3], [3, 4, 4]]
print(data.minimumCost(n, connections))
